from flask import Flask, request, jsonify
import jieba.analyse
import jieba.posseg as pseg

app = Flask(__name__)


@app.route('/health', methods=['GET'])
def health():
    return "ok", 200


@app.route('/extract_keywords', methods=['POST'])
def extract_keywords():
    try:
        # 解析请求数据
        data = request.get_json()
        text = data.get('text')
        top_k = data.get('top_k', 5)  # 默认返回5个关键词

        # 验证文本长度（1万字=10000字符）
        if not text:
            return jsonify({'error': '文本内容不能为空'}), 400
        if len(text) > 10000:
            return jsonify({'error': '文本长度超过1万字限制'}), 400

        # 验证关键词数量参数
        if not isinstance(top_k, int) or top_k < 1:
            return jsonify({'error': 'top_k必须为正整数'}), 400

        # 使用jieba提取关键词（TF-IDF算法）
        keywords = jieba.analyse.extract_tags(
            sentence=text,
            topK=top_k,
            withWeight=False,
            allowPOS=()  # 不限制词性
        )

        return jsonify({
            'success': True,
            'keywords': keywords,
            'count': len(keywords)
        })

    except Exception as e:
        return jsonify({'error': f'服务器内部错误: {str(e)}'}), 500


@app.route('/extract_entities', methods=['POST'])
def extract_entities():
    try:
        # 解析请求数据
        data = request.get_json()
        text = data.get('text')

        # 验证文本长度（1万字=10000字符）
        if not text:
            return jsonify({'error': '文本内容不能为空'}), 400
        if len(text) > 10000:
            return jsonify({'error': '文本长度超过1万字限制'}), 400

        # 使用jieba提取实体（依据词性）
        words = pseg.cut(text)

        # 提取人名、地名、组织名
        person_names = []
        locations = []
        organizations = []

        for word, flag in words:
            if len(word) <= 1:
                continue
            if flag == 'nr':  # 人名
                person_names.append(word)
            elif flag == 'ns':  # 地名
                locations.append(word)
            elif flag == 'nt':  # 组织名
                organizations.append(word)

        return jsonify({
            'success': True,
            'keywords': {
                'persons': list(set(person_names)),
                'locations': list(set(locations)),
                'organizations': list(set(organizations))
            }
        })

    except Exception as e:
        return jsonify({'error': f'服务器内部错误: {str(e)}'}), 500


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)
